import argparse
import os
import asyncio
import json

from dataloader.df_dataloader import load_dataset
from util.ark_r1_api_batch import ark_batch_query
from util.batch_generate_api import batch_generate


def save_json(data, json_path):
    with open(json_path, "w") as f:
        json.dump(data, f, ensure_ascii=False, indent=4)


def get_vllm_results(dataset, base_url, api_key, model, worker_num, system_prompt):
    dataset = asyncio.run(
        batch_generate(
            base_url,
            api_key,
            model,
            dataset,
            concurrent=worker_num,
            system_prompt=system_prompt,
        )
    )
    return dataset


def get_ark_results(model, ark_model_id, dataset, worker_num):
    dataset = asyncio.run(
        ark_batch_query(
            model, ark_model_id, dataset, concurrent=worker_num, is_force_think=False
        )
    )
    return dataset


def infer(
    model_name, ark_r1_batch_model_id, base_url, api_key, dataset, worker_num, save_path
):
    """大模型推理获得结果"""
    if model_name == "deepseek-r1":
        # 大模型推理
        dataset = get_ark_results(
            ark_r1_batch_model_id, model_name, dataset, worker_num
        )
    else:
        # 防止模型不思考
        cot_start_tag = "<think>"
        cot_end_tag = "</think>"
        system_prompt = f"任何输出都要有思考过程，思考过程必须以“{cot_start_tag}\n\n嗯”开头，以“\n{cot_end_tag}\n\n”结尾，你要仔细揣摩用户意图。在思考过程后，紧接着提供逻辑清晰且内容完整的回答。\n你是一名熟悉中国期货市场的金融专家, 充分了解期货交易规则、期货品种、期货投资策略、期货期权的定价、业务合规等方面的知识。请你回答以下期货相关问题。"
        dataset = get_vllm_results(
            dataset, base_url, api_key, model_name, worker_num, system_prompt
        )
    os.makedirs(save_path, exist_ok=True)
    # 保存为 csv
    try:
        save_json(dataset, os.path.join(save_path, f"{model_name}_ga.json"))
    except Exception as e:
        print(dataset)
        print(f"Error occurred while saving as json: {e}")


def fineva_main(args):
    test_file = args.test_file
    model_name = args.model_name
    save_path = args.save_path
    base_url = args.base_url
    api_key = args.api_key
    worker_num = args.worker_num
    sid_list = args.sid_list
    ark_r1_batch_model_id = args.ark_r1_batch_model_id
    # 导入数据
    dataset = load_dataset(test_file, sid_list)
    infer(
        model_name,
        ark_r1_batch_model_id,
        base_url,
        api_key,
        dataset,
        worker_num,
        save_path,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name", default="deepseek-r1", type=str, help="vllm API支持的模型名称"
    )
    parser.add_argument("--model_path", required=False, type=str)
    parser.add_argument(
        "--base_url",
        default="http://127.0.0.1:8000/v1",
        type=str,
        help="vllm api服务器的url，默认8000端口",
    )
    parser.add_argument(
        "--api_key", default="pobo2025", type=str, help="vllm API 的 api-key"
    )
    parser.add_argument(
        "--worker_num",
        required=False,
        type=int,
        default=20,
        help="vllm API 并发数，根据vllm服务器的显卡数量决定",
    )
    parser.add_argument(
        "--test_file",
        default="data/DFBenchmark.json",
        type=str,
        help="测试集路径",
    )
    parser.add_argument("--save_path", default="result", type=str)
    parser.add_argument(
        "--sid_list", default="all", type=str, help="评测维度列表，用英文逗号分割"
    )
    parser.add_argument(
        "--ark_r1_batch_model_id",
        required=True,
        type=str,
        help="火山引擎 DS R1 模型id",
    )
    args = parser.parse_args()
    fineva_main(args)
